#!/usr/bin/env python
# author gesturewei, ihep, cas


def usage():
    print("Usage: check_mem -w <percent_free>% -c <percent_free>%");
    print("\t -w <bytes_free> -c <bytes_free>");


def main():
    EXIT_OK=0;
    EXIT_WARNING=1;
    EXIT_CRITICAL=2;
    EXIT_UNKNOWN=3;
    import sys, psutil, getopt;
    warning_value_type=0; # 0 not set; 1 percent free set; 2 bytes free set
    critical_value_type=0; # 0 not set; 1 percent free set; 2 bytes free set
    warning_value=0.0;
    critical_value=0.0;
    try:
        opts, args=getopt.getopt(sys.argv[1:],"w:c:h");
    except getopt.GetoptError:
        usage();
        sys.exit(EXIT_UNKNOWN);
    for key,value in opts:
        if (key=='-w'):
            if (value.endswith('%')):
                warning_value_type=1;
                warning_value=float(value[:-1]);
            else:
                warning_value_type=2;
                warning_value=float(value);
        elif (key=='-c'):
            if (value.endswith('%')):
                critical_value_type=1;
                critical_value=float(value[:-1]);
            else:
                critical_value_type=2;
                critical_value=float(value);
        elif (key=='-h'):
            usage();
            sys.exit(EXIT_OK);
    if (warning_value_type==0 or critical_value_type==0):
        usage();
        sys.exit(EXIT_UNKNOWN);
    if (warning_value<0.0 or critical_value<0.0):
        print("The value must be positive.");
        sys.exit(EXIT_UNKNOWN);
    mem_info=psutil.virtual_memory();
    if (warning_value_type==1):
        warning_value*=mem_info[0]/100;
    if (critical_value_type==1):
        critical_value*=mem_info[0]/100;
    if (mem_info[1]<=critical_value):
        out_msg="MEMORY CRITICAL - ";
        exit_code=EXIT_CRITICAL;
    elif(mem_info[1]<=warning_value):
        out_msg="MEMORY WARNING - ";
        exit_code=EXIT_WARNING;
    else:
        out_msg="MEMORY OK - ";
        exit_code=EXIT_OK;
    total_in_mb=mem_info[0]>>20;
    available_in_mb=mem_info[1]>>20;
    out_msg+="%d%% free (%d MB out of %d MB) |mem=" % (100-mem_info[2],available_in_mb,total_in_mb);
    out_msg+="%dMB;%d;%d;0;%d" % (available_in_mb,int(warning_value/1048576),int(critical_value/1048576),total_in_mb);
    print(out_msg);
    sys.exit(exit_code);

if __name__=="__main__":
    main();